import re

import chromadb
from chromadb.utils import embedding_functions
from docx import Document
import os
from openai import OpenAI
# 添加PyPDFLoader导入
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter


def load_recipe_file(file_path):
    """
    使用PyPDFLoader加载菜谱PDF文件

    参数:
        file_path: PDF文件路径

    返回:
        文档内容字符串
    """
    print(f"正在使用PyPDFLoader加载文件: {file_path}")

    # 检查文件是否存在
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"文件不存在: {file_path}")

    # 使用PyPDFLoader加载PDF文件
    loader = PyPDFLoader(file_path)
    documents = loader.load()

    # 合并所有页面的内容
    full_text = ""
    for doc in documents:
        full_text += doc.page_content + "\n"

    print(f"成功加载PDF文件，共{len(documents)}页，文本长度: {len(full_text)}字符")
    return full_text


def split_text(text, chunk_size=200):
    # 使用RecursiveCharacterTextSplitter替代简单的文本分割
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=20,  # 添加重叠以保持上下文连贯性
        length_function=len,
        separators=["\n\n", "\n", ".", "!", "?", ":", ";", ",", " ", ""]
    )

    # 分割文本
    chunks = text_splitter.split_text(text)
    print(f"文本已分割为 {len(chunks)} 个块")
    return chunks


def text_embeddings(texts):
    client = OpenAI(
        api_key="ollama",  # Ollama不需要API密钥，但openai库需要一个值
        base_url="http://localhost:11434/v1"  # Ollama服务的base_url
    )

    # 如果输入是字符串，转换为列表
    if isinstance(texts, str):
        texts = [texts]

    # Ollama的嵌入模型批处理大小可以更大，但我们保持一致
    batch_size = 32
    embeddings = []

    # 分批处理输入文本
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        completion = client.embeddings.create(
            model="nomic-embed-text:latest",  # Ollama的嵌入模型
            input=batch,
            encoding_format="float"
        )

        print(f"Processing batch {i + 1}:{completion}")

        # 提取嵌入向量
        batch_embeddings = [item.embedding for item in completion.data]
        embeddings.extend(batch_embeddings)

    return embeddings


def store_in_chromadb(chunks, collection_name, persist_directory="./chroma_db"):
    # 1. 生成文本嵌入向量
    embeddings = text_embeddings(chunks)
    if len(embeddings) != len(chunks):
        print("嵌入向量生成异常，数量与文本块不匹配")
        return

    # 2. 初始化ChromaDB客户端并创建/获取集合
    client = chromadb.PersistentClient(path=persist_directory)
    # 创建集合时指定嵌入函数（此处使用已生成的嵌入，故设为None）
    collection = client.get_or_create_collection(
        name=collection_name,
        embedding_function=None  # 不使用内置嵌入函数，使用自定义生成的嵌入
    )

    # 3. 生成唯一ID（格式：chunk_0, chunk_1, ...）
    ids = [f"chunk_{i}" for i in range(len(chunks))]

    # 4. 将数据存入集合
    collection.add(
        documents=chunks,  # 文本片段
        embeddings=embeddings,  # 对应的嵌入向量
        ids=ids  # 唯一标识
    )

    print(f"成功存入ChromaDB集合 '{collection_name}'，共 {len(chunks)} 个文本块")


def query_chroma(query_text):
    # 创建嵌入函数
    embedding_function = embedding_functions.OpenAIEmbeddingFunction(
        api_key="ollama",
        model_name="nomic-embed-text:latest",
        api_base="http://localhost:11434/v1"
    )

    # 2. 初始化ChromaDB客户端并创建/获取集合
    persist_directory = "./chroma_db"
    chroma_client = chromadb.PersistentClient(path=persist_directory)
    # 获取集合时指定嵌入函数
    collection = chroma_client.get_collection(
        name="recipe_collection",  # 修改为正确的集合名称
        embedding_function=embedding_function
    )

    # 查询与特定文本相似的向量
    results = collection.query(
        query_texts=[query_text],
        n_results=5,  # 返回最相似的5个结果
        include=['embeddings', 'documents', 'metadatas', 'distances']  # 显式指定需要返回的字段
    )
    # {'ids': [['text_41', 'text_33', 'text_37', 'text_35', 'text_44']], 'embeddings': None, 'documents': [['年假（5-15天，按工作年限）、病假（按国家规定）、婚假/产假/陪产假等法定假期\n健康福利：每年1次', '到岗\n相应调整下班时间，保证8小时工作时长\n需要提前向直属主管申请并获得批准\n特殊工作安排\n项目紧急', '分钟以上\n早退：提前离开工作岗位\n旷工：未经批准不到岗工作\n考勤处理\n迟到3次以内：口头提醒\n迟到3', '要求\n员工应按时上下班，严格遵守考勤制度\n使用公司考勤系统进行打卡签到\n因故不能按时到岗应提前请假\n', '、年度旅游/聚餐、员工家庭日、兴趣社团支持\n第六章 培训发展\n6.1 培训体系\n新员工培训（3周）\n']], 'uris': None, 'included': ['metadatas', 'documents', 'distances'], 'data': None, 'metadatas': [[None, None, None, None, None]], 'distances': [[0.6660666465759277, 0.8344534039497375, 0.9031039476394653, 0.9283781051635742, 0.9392238855361938]]}

    print(results)
    return results


def rag_answer_question(results, question):
    retrieved_chunks = list(set(results["documents"][0]))
    print("retrieved_chunks", retrieved_chunks)

    # 构建上下文（把片段拼起来，告诉LLM"基于这些内容回答"）
    context = "\n".join([f"- {chunk}" for chunk in retrieved_chunks])
    print(f"\n【检索到的参考内容】\n{context}\n")
    prompt = f"""
你是一个小红书百科小助手，记录菜谱、手工制作等知识点，请基于提供的文档内容回答用户问题。

【参考内容】：
{context}

【用户问题】：
{question}

【回答要求】：
1. 严格基于参考内容回答，不要编造或推测
2. 如果参考内容中没有相关信息，请明确说明"参考内容中未提及此信息"
3. 回答应简洁明了，条理清晰
4. 如有多个要点，请分点列出
5. 如果涉及具体制作步骤，可引用相关内容
6. 回答完成后，可以提供一些相关的制作建议
    """

    print("提示词:", prompt)

    messages = [{"role": "user", "content": prompt}]
    ai_response = get_response(messages)
    print("大模型返回的答案:", ai_response)


def get_response(messages):
    client = OpenAI(
        # 使用Ollama本地部署的Qwen3模型
        api_key="ollama",  # Ollama不需要API密钥，但openai库需要一个值
        base_url="http://localhost:11434/v1",  # Ollama服务的base_url
    )

    completion = client.chat.completions.create(
        model="qwen3",  # 使用Ollama本地部署的Qwen3模型
        messages=messages,
    )
    return completion.choices[0].message.content


if __name__ == '__main__':
    # 使用PyPDFLoader加载菜谱PDF文件
    try:
        # 构建菜谱文件路径
        recipe_file_path = os.path.join(os.path.dirname(__file__), "菜谱.pdf")
        print(f"菜谱文件路径: {recipe_file_path}")

        # 加载菜谱文件
        recipe_text = load_recipe_file(recipe_file_path)
        print("菜谱内容预览:")
        print(recipe_text[:500] + "..." if len(recipe_text) > 500 else recipe_text)

        # 分割文本
        chunks = split_text(recipe_text)
        print(f"生成的文本块数量: {len(chunks)}")
        print("第一个文本块预览:")
        print(chunks[0][:200] + "..." if len(chunks[0]) > 200 else chunks[0])

        # 存储到ChromaDB
        store_in_chromadb(chunks, "recipe_collection")

        # 测试查询
        question = "拔丝橘子怎么做？"
        results = query_chroma(question)
        rag_answer_question(results, question)

    except FileNotFoundError as e:
        print(f"文件未找到错误: {e}")
        print("请确保菜谱.pdf文件存在于views目录中")
    except Exception as e:
        print(f"处理菜谱文件时出错: {e}")
        import traceback

        traceback.print_exc()